import time
import logging
from typing import List, Dict, Any, Optional
from io import BytesIO
import numpy as np
from PIL import Image
import cv2

try:
    from paddleocr import PaddleOCR
except ImportError:
    PaddleOCR = None

from config.settings import settings
from models.response_models import OCRResult, OCRData
from utils.preprocessing_utils import preprocess_image_for_ocr

logger = logging.getLogger(__name__)

class OCRService:
    """OCR service using PaddleOCR"""
    
    def __init__(self):
        self.ocr = None
        self.is_initialized = False
        
    def initialize(self):
        """Initialize PaddleOCR with configurable model options"""
        try:
            if PaddleOCR is None:
                raise ImportError("PaddleOCR is not installed. Please install it using: pip install paddleocr")
            
            logger.info("Initializing PaddleOCR...")
            
            # Prepare OCR initialization parameters
            ocr_params = {
                'use_angle_cls': settings.OCR_USE_ANGLE_CLS,
                'lang': settings.OCR_LANGUAGE
            }
            
            # Add model directory parameters if specified
            if settings.OCR_DET_MODEL_DIR:
                ocr_params['det_model_dir'] = settings.OCR_DET_MODEL_DIR
            elif settings.OCR_USE_HIGH_ACCURACY:
                # Use server detection model for higher accuracy
                ocr_params['det_model_dir'] = None  # Let PaddleOCR use default server model
            
            if settings.OCR_REC_MODEL_DIR:
                ocr_params['rec_model_dir'] = settings.OCR_REC_MODEL_DIR
            elif settings.OCR_USE_HIGH_ACCURACY:
                # Use server recognition model for higher accuracy
                ocr_params['rec_model_dir'] = None  # Let PaddleOCR use default server model
                
            if settings.OCR_CLS_MODEL_DIR:
                ocr_params['cls_model_dir'] = settings.OCR_CLS_MODEL_DIR
            
            logger.info(f"PaddleOCR initialization parameters: {ocr_params}")
            self.ocr = PaddleOCR(**ocr_params)
            self.is_initialized = True
            logger.info("PaddleOCR initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize PaddleOCR: {e}")
            raise RuntimeError(f"OCR initialization failed: {e}")
    
    def is_healthy(self) -> bool:
        """Check if OCR service is healthy"""
        return self.is_initialized and self.ocr is not None
    
    def _preprocess_image(self, image_data: bytes, enhance: bool = False, 
                         correct_skew: bool = False) -> np.ndarray:
        """Preprocess image data for OCR with optional enhancement and skew correction"""
        try:
            # Apply preprocessing if requested
            if enhance or correct_skew:
                processed_data = preprocess_image_for_ocr(
                    image_data, 
                    enhance=enhance, 
                    correct_skew_flag=correct_skew
                )
            else:
                processed_data = image_data
            
            # Load image from bytes
            image = Image.open(BytesIO(processed_data))
            
            # Convert to RGB if necessary
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Convert PIL image to numpy array
            img_array = np.array(image)
            
            # Convert RGB to BGR for OpenCV
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
            
            return img_array
            
        except Exception as e:
            logger.error(f"Image preprocessing failed: {e}")
            raise ValueError(f"Failed to process image: {e}")
    
    def recognize_image(self, image_data: bytes, language: str = "ch", confidence_threshold: float = 0.5,
                       enhance: bool = False, correct_skew: bool = False) -> OCRData:
        """Recognize text from image data with optional preprocessing
        
        Args:
            image_data: Image data as bytes
            language: OCR language setting
            confidence_threshold: Minimum confidence threshold for detection
            enhance: Whether to apply image enhancement
            correct_skew: Whether to apply skew correction
        """
        if not self.is_healthy():
            raise RuntimeError("OCR service is not initialized")
        
        start_time = time.time()
        
        try:
            # Preprocess image
            img_array = self._preprocess_image(image_data, enhance, correct_skew)
            
            # Perform OCR
            logger.info(f"Performing OCR recognition with language: {language}...")
            ocr_results = self.ocr.ocr(img_array)
            
            # Process results
            results = []
            all_text = []
            confidences = []
            
            if ocr_results and ocr_results[0]:
                for line in ocr_results[0]:
                    if line:
                        bbox = line[0]  # Bounding box coordinates
                        text_info = line[1]  # (text, confidence)
                        
                        if text_info and len(text_info) >= 2:
                            text = text_info[0]
                            confidence = float(text_info[1])
                            
                            # Filter by confidence threshold
                            if confidence >= confidence_threshold:
                                # Convert bbox to list of integers
                                bbox_int = [[int(coord[0]), int(coord[1])] for coord in bbox]
                                
                                results.append(OCRResult(
                                    text=text,
                                    confidence=confidence,
                                    bbox=bbox_int
                                ))
                                
                                all_text.append(text)
                                confidences.append(confidence)
            
            processing_time = time.time() - start_time
            total_text = ' '.join(all_text)
            total_characters = len(total_text.replace(' ', ''))
            average_confidence = sum(confidences) / len(confidences) if confidences else 0.0
            
            logger.info(f"OCR completed in {processing_time:.2f}s, found {len(results)} text regions")
            
            return OCRData(
                results=results,
                total_text=total_text,
                processing_time=processing_time,
                language_used=language,
                total_characters=total_characters,
                average_confidence=average_confidence
            )
            
        except Exception as e:
            logger.error(f"OCR recognition failed: {e}")
            raise RuntimeError(f"OCR recognition failed: {e}")
    
    def recognize_digits_only(self, image_data: bytes, confidence_threshold: float = 0.8, 
                              filter_spaces: bool = True, filter_special_chars: bool = True,
                              enhance: bool = False, correct_skew: bool = False) -> OCRData:
        """Recognize only digits from image data, filtering out spaces and special characters
        
        Args:
            image_data: Image data as bytes
            confidence_threshold: Minimum confidence threshold for digit detection (0.0-1.0)
            filter_spaces: Whether to filter out spaces from results
            filter_special_chars: Whether to filter out special characters
            enhance: Whether to apply image enhancement
            correct_skew: Whether to apply skew correction
        """
        if not self.is_healthy():
            raise RuntimeError("OCR service is not initialized")
        
        start_time = time.time()
        
        try:
            # Preprocess image
            img_array = self._preprocess_image(image_data, enhance, correct_skew)
            
            # Perform OCR
            logger.info("Performing digit-only OCR recognition...")
            ocr_results = self.ocr.ocr(img_array)
            
            # Process results
            results = []
            all_text = []
            confidences = []
            
            if ocr_results and ocr_results[0]:
                for line in ocr_results[0]:
                    if line:
                        bbox = line[0]  # Bounding box coordinates
                        text_info = line[1]  # (text, confidence)
                        
                        if text_info and len(text_info) >= 2:
                            text = text_info[0]
                            confidence = float(text_info[1])
                            
                            # Filter by confidence threshold
                            if confidence >= confidence_threshold:
                                # Filter text to digits only if requested
                                if filter_spaces or filter_special_chars:
                                    # Keep only digits (0-9)
                                    filtered_text = ''.join(char for char in text if char.isdigit())
                                else:
                                    filtered_text = text
                                
                                # Only add results that contain digits
                                if filtered_text:
                                    # Convert bbox to list of integers
                                    bbox_int = [[int(coord[0]), int(coord[1])] for coord in bbox]
                                    
                                    results.append(OCRResult(
                                        text=filtered_text,
                                        confidence=confidence,
                                        bbox=bbox_int
                                    ))
                                    
                                    all_text.append(filtered_text)
                                    confidences.append(confidence)
            
            processing_time = time.time() - start_time
            total_text = ''.join(all_text) if filter_spaces else ' '.join(all_text)
            total_characters = len(total_text.replace(' ', ''))
            average_confidence = sum(confidences) / len(confidences) if confidences else 0.0
            
            logger.info(f"Digit-only OCR completed in {processing_time:.2f}s, found {len(results)} digit regions")
            
            return OCRData(
                results=results,
                total_text=total_text,
                processing_time=processing_time,
                language_used="digits",
                total_characters=total_characters,
                average_confidence=average_confidence
            )
            
        except Exception as e:
            logger.error(f"Digit-only OCR recognition failed: {e}")
            raise RuntimeError(f"Digit-only OCR recognition failed: {e}")